import torch
import torch.nn as nn

class BranchNet(nn.Module):
    """独立的四层全连接分支网络"""
    def __init__(self, input_dim):
        super().__init__()
        # 定义四层全连接层（3个隐藏层 + 1个输出层）
        self.fc1 = nn.Linear(input_dim, 1000)
        self.fc2 = nn.Linear(1000, 1000)
        self.fc3 = nn.Linear(1000, 1000)
        self.fc4 = nn.Linear(1000, input_dim)  # 输出层恢复原始展平维度

    def forward(self, x, index):
        # 前三个全连接层后接sigmoid零中心化
        x = torch.relu(self.fc1(x))
        x = torch.relu(self.fc2(x))
        x = torch.relu(self.fc3(x))
        if index == 1:
            x = torch.sigmoid(self.fc4(x))  # 左子网络
        else:
            x = torch.sigmoid(self.fc4(x)) - 1 #右子网络
        return x

class ParallelImageNet(nn.Module):
    """双分支并行图像处理网络"""
    def __init__(self, batch_size, image_size, in_channels=3):
        
        super().__init__()
        self.batch_size = batch_size
        self.image_size = image_size
        self.in_channels = in_channels
        self.flatten_dim = in_channels * image_size * image_size  # 展平后的输入维度
        
        # 创建两个独立分支网络
        self.branch_a = BranchNet(self.flatten_dim).cuda()
        self.branch_b = BranchNet(self.flatten_dim).cuda()

    def forward(self, x):
        batch_size = self.batch_size
        channels = self.in_channels
        image_size = self.image_size
        # 生成测试数据：[batch, 3, 128, 128]
        dummy_input = torch.rand(batch_size, channels, image_size, image_size)
        # 输入x形状: [batch_size, image_size, image_size]
        dummy_input = dummy_input.cuda()  # 关键：重新赋值到原变量
        # Step 1: 展平输入图像
        dummy_input_flat = dummy_input.view(batch_size, -1)  # 变为[batch_size, flatten_dim]
        
        # Step 2: 双分支并行处理
        out_a = self.branch_a(dummy_input_flat, index=1)
        out_b = self.branch_b(dummy_input_flat, index=2)
        
        # Step 3: 结果按位相加并恢复图像形状
        combined = out_a + out_b
        return x + combined.view(batch_size, self.in_channels, 
                            self.image_size, self.image_size)

# 示例用法（RGB三通道）
if __name__ == "__main__":
    # 配置参数
    image_size = 128
    channels = 3
    batch_size = 2
    
    # 生成测试数据：[batch, 3, 128, 128]
    dummy_input = torch.rand(batch_size, channels, image_size, image_size)
    print("输入形状:", dummy_input.shape)  # 期望输出: torch.Size([2, 3, 128, 128])
    dummy_input = dummy_input.cuda()
    # 初始化模型
    model = ParallelImageNet(batch_size, image_size, in_channels=channels).cuda()
    
    # 前向传播
    output = model(dummy_input)
    print("\n输出形状:", output.shape)  # 应保持与输入相同形状
    print("输出范围:", output.min().item(), "~", output.max().item())  # 目标范围[-1, 1]